package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"
)

type Config struct {
	ImageName      string
	ImageTag       string
	TarFileName    string
	RemoteHost     string
	RemotePort     string
	RemoteUser     string
	RemotePassword string
	Registry       string
	RegistryUser   string
	RegistryPass   string
	ShowHelp       bool
	SkipPush       bool
}

func main() {
	config := parseFlags()

	if config.ShowHelp {
		printUsage()
		return
	}

	if config.ImageName == "" {
		fmt.Println("错误: 必须指定镜像名称")
		printUsage()
		os.Exit(1)
	}

	fmt.Println("*** 开始镜像处理流程 ***")

	// 1. 拉取原始镜像
	fmt.Println("*** 正在拉取镜像...")
	if err := pullImage(config.ImageName); err != nil {
		log.Fatalf("错误: 拉取镜像失败: %v", err)
	}

	// 2. 打标签
	fmt.Println("*** 正在标记镜像...")
	taggedImage := getTaggedImageName(config.ImageName, config.ImageTag)
	if err := tagImage(config.ImageName, taggedImage); err != nil {
		log.Fatalf("错误: 标记镜像失败: %v", err)
	}

	// 3. 保存为tar包
	fmt.Println("*** 正在保存镜像...")
	if err := saveImage(taggedImage, config.TarFileName); err != nil {
		log.Fatalf("错误: 保存镜像失败: %v", err)
	}

	// 4. 传输tar包到远程主机
	fmt.Println("*** 正在传输文件到远程主机...")
	if err := transferFile(config.TarFileName, config.RemoteHost, config.RemotePort,
		config.RemoteUser, config.RemotePassword); err != nil {
		log.Fatalf("错误: 文件传输失败: %v", err)
	}

	// 5. 在远程主机上加载镜像并推送到仓库
	fmt.Println("*** 正在远程处理镜像...")
	if err := remoteLoadAndPush(config, taggedImage); err != nil {
		log.Fatalf("错误: 远程操作失败: %v", err)
	}

	// 6. 清理本地文件
	fmt.Println("*** 正在清理临时文件...")
	if err := cleanup(taggedImage, config.TarFileName); err != nil {
		log.Printf("警告: 清理失败: %v", err)
	}
	if err := removeAllLocalImages(); err != nil {
		log.Printf("警告: 清理本地镜像失败: %v", err)
	}

	fmt.Println("*** 操作成功完成! ***")
}

func remoteExec(host, port, user, password, command string) error {
	// 构造SSH命令
	sshCmd := exec.Command("ssh", "-p", port,
		"-o", "StrictHostKeyChecking=no",
		"-o", "UserKnownHostsFile=/dev/null",
		fmt.Sprintf("%s@%s", user, host), command)

	// 获取输入管道
	stdin, err := sshCmd.StdinPipe()
	if err != nil {
		return fmt.Errorf("创建输入管道失败: %v", err)
	}

	// 设置输出
	sshCmd.Stdout = os.Stdout
	sshCmd.Stderr = os.Stderr

	// 启动命令
	if err := sshCmd.Start(); err != nil {
		return fmt.Errorf("启动SSH命令失败: %v", err)
	}

	// 等待密码提示并自动输入
	time.Sleep(1 * time.Second) // 等待密码提示出现
	if _, err := stdin.Write([]byte(password + "\n")); err != nil {
		return fmt.Errorf("写入密码失败: %v", err)
	}
	stdin.Close()

	// 等待命令完成
	return sshCmd.Wait()
}

func transferFile(fileName, host, port, user, password string) error {
	// 构造SCP命令
	scpCmd := exec.Command("scp", "-P", port,
		"-o", "StrictHostKeyChecking=no",
		"-o", "UserKnownHostsFile=/dev/null",
		fileName, fmt.Sprintf("%s@%s:~/%s", user, host, fileName))

	// 获取输入管道
	stdin, err := scpCmd.StdinPipe()
	if err != nil {
		return fmt.Errorf("创建输入管道失败: %v", err)
	}

	// 设置输出
	scpCmd.Stdout = os.Stdout
	scpCmd.Stderr = os.Stderr

	// 启动命令
	if err := scpCmd.Start(); err != nil {
		return fmt.Errorf("启动SCP命令失败: %v", err)
	}

	// 等待密码提示并自动输入
	time.Sleep(1 * time.Second) // 等待密码提示出现
	if _, err := stdin.Write([]byte(password + "\n")); err != nil {
		return fmt.Errorf("写入密码失败: %v", err)
	}
	stdin.Close()

	// 等待命令完成
	return scpCmd.Wait()
}

func remoteLoadAndPush(config Config, taggedImage string) error {
	// 1. 在远程主机上加载镜像
	loadCmd := fmt.Sprintf("docker load -i ~/%s", config.TarFileName)
	if err := remoteExec(config.RemoteHost, config.RemotePort, config.RemoteUser,
		config.RemotePassword, loadCmd); err != nil {
		return fmt.Errorf("远程加载镜像失败: %v", err)
	}

	// 2. 推送镜像
	pushCmd := fmt.Sprintf("docker push %s", taggedImage)
	if err := remoteExec(config.RemoteHost, config.RemotePort, config.RemoteUser,
		config.RemotePassword, pushCmd); err != nil {
		return fmt.Errorf("推送镜像失败: %v", err)
	}

	// 3. 清理远程主机上的文件
	cleanupCmd := fmt.Sprintf("rm -f ~/%s && docker rmi %s",
		config.TarFileName, taggedImage)
	if err := remoteExec(config.RemoteHost, config.RemotePort, config.RemoteUser,
		config.RemotePassword, cleanupCmd); err != nil {
		log.Printf("警告: 远程清理失败: %v", err)
	}

	return nil
}

func parseFlags() Config {
	var config Config

	flag.StringVar(&config.ImageName, "image", "", "要拉取的镜像名称(必填)")
	flag.StringVar(&config.ImageTag, "tag", "", "镜像标签(默认: harbor.linux.com/cicd_registory/<镜像>:<版本>)")
	flag.StringVar(&config.TarFileName, "tar", "image.tar", "保存的压缩包文件名")
	flag.StringVar(&config.RemoteHost, "host", "192.168.3.82", "远程主机地址")
	flag.StringVar(&config.RemotePort, "port", "10022", "远程SSH端口")
	flag.StringVar(&config.RemoteUser, "user", "root", "远程用户名(默认: root)")
	flag.StringVar(&config.RemotePassword, "password", "666666", "远程用户密码(默认: 666666)")
	flag.StringVar(&config.Registry, "registry", "", "Docker仓库地址(如 harbor.linux.com)")
	flag.StringVar(&config.RegistryUser, "registry-user", "", "仓库用户名")
	flag.StringVar(&config.RegistryPass, "registry-pass", "", "仓库密码")
	flag.BoolVar(&config.SkipPush, "skip-push", false, "跳过推送步骤")
	flag.BoolVar(&config.ShowHelp, "help", false, "显示帮助信息")

	flag.Parse()

	return config
}

func printUsage() {
	fmt.Println(`用法: docker-image-processor -image <镜像名称> [选项]

镜像处理工具：拉取镜像 → 标记 → 打包 → 传输 → 远程加载

必选参数:
  -image string    要拉取的镜像名称 (如 "nginx:1.21")

常用选项:
  -tag string      自定义标记名称 (默认: "harbor.linux.com/cicd_registory/<镜像>:<版本>")
  -tar string      保存的压缩包文件名 (默认: "image.tar")
  -host string     远程主机地址 (默认: "192.168.3.82")
  -port string     远程SSH端口 (默认: "10022")
  -user string     远程用户名 (默认: "root")
  -password string 远程用户密码 (默认: "666666")
  -help            显示帮助信息

高级选项:
  -registry string       Docker仓库地址
  -registry-user string  仓库用户名
  -registry-pass string  仓库密码
  -skip-push             跳过推送步骤

示例:
  1. 基本用法:
     docker-image-processor -image nginx:1.21

  2. 自定义标记和文件名:
     docker-image-processor -image nginx:1.21 -tag myregistry.com/myproject/nginx:1.21 -tar nginx.tar

  3. 自定义远程主机:
     docker-image-processor -image nginx:1.21 -host 192.168.1.100 -port 22 -user admin -password 123456`)
}

func pullImage(imageName string) error {
	cmd := exec.Command("docker", "pull", "--platform", "linux/amd64", imageName)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func getTaggedImageName(originalImage, customTag string) string {
	if customTag != "" {
		return customTag
	}

	parts := strings.Split(originalImage, ":")
	imageName := parts[0]
	var version string
	if len(parts) > 1 {
		version = parts[1]
	} else {
		version = "latest"
	}

	baseNameParts := strings.Split(imageName, "/")
	baseName := baseNameParts[len(baseNameParts)-1]

	return fmt.Sprintf("harbor.linux.com/cicd_registory/%s:%s", baseName, version)
}

func tagImage(originalImage, taggedImage string) error {
	cmd := exec.Command("docker", "tag", originalImage, taggedImage)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func saveImage(imageName, tarFileName string) error {
	cmd := exec.Command("docker", "save", "-o", tarFileName, imageName)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func cleanup(imageName, tarFileName string) error {
	if err := exec.Command("docker", "rmi", "-f", imageName).Run(); err != nil {
		return fmt.Errorf("删除镜像失败 %s: %v", imageName, err)
	}

	if err := os.Remove(tarFileName); err != nil {
		return fmt.Errorf("删除文件失败 %s: %v", tarFileName, err)
	}

	return nil
}

func removeAllLocalImages() error {
	// 获取所有镜像ID
	listCmd := exec.Command("docker", "image", "ls", "-q")
	output, err := listCmd.Output()
	if err != nil {
		return fmt.Errorf("获取镜像列表失败: %v", err)
	}

	imageIDs := strings.Fields(string(output))
	if len(imageIDs) == 0 {
		return nil
	}

	// 删除所有镜像
	rmCmd := exec.Command("docker", "rmi", "-f")
	rmCmd.Args = append(rmCmd.Args, imageIDs...)
	rmCmd.Stdout = os.Stdout
	rmCmd.Stderr = os.Stderr

	if err := rmCmd.Run(); err != nil {
		return fmt.Errorf("删除镜像失败: %v", err)
	}

	return nil
}
